'''
Created on Nov 23, 2010

@author: oabalbin
'''
import os
import exome.gatk_cluster.picard_commands_local as mypicard
import logging
from exome.jobs.base import JOB_SUCCESS, JOB_ERROR, up_to_date,job_done,up_to_date_job
from optparse import OptionParser
from exome.jobs.config import ExomePipelineConfig, ExomeLaneConfig


def count_covariates(reference_genome, dbsnp_ref, bam_file_name, outfile, use_mem, num_cores, path_to_gatk):
    '''
    Generated the file of count covariates, which is required for base quality recalibration
    
    java -Xmx16g -jar /exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/GenomeAnalysisTK.jar 
    -l INFO -nt 6 -R /exds/projects/alignment_indexes/gatk/hg19/hg19.fa 
    -B:mask,VCF /exds/projects/alignment_indexes/gatk/hg19/dbsnp132_00-All_processed5.processed.vcf 
    -I s_3_12_sequence.hg19.aln2.rmdup.sorted.bam -T CountCovariates -cov ReadGroupCovariate 
    -cov QualityScoreCovariate -cov CycleCovariate -cov DinucCovariate 
    -recalFile s_3_12_sequence.hg19.aln2.rmdup.sorted.recal_data.csv
    
    
    Note: you can use the --process_nth_locus/ -pN 3 to speed up the process. 
    Which force the program to evaluate every 3 loci, which makes to skip in general ~66 % data
    when dbsnp is also included. Important for Large BAM files
    '''
    
    gatk_command=path_to_gatk+'GenomeAnalysisTK.jar'
    #outfile=bam_file_name.replace('.bam','.recal_data.csv')
    args = ['java','-Xmx'+str(use_mem)+'m', '-jar',gatk_command,
            '-l', 'INFO', '-nt', str(num_cores), '-R',reference_genome, 
            '-B:dbsnp,VCF', dbsnp_ref, '-I', bam_file_name, '-T','CountCovariates',
            '-cov', 'ReadGroupCovariate','-cov', 'QualityScoreCovariate', '-cov', 
            'CycleCovariate', '-cov', 'DinucCovariate', '-recalFile',outfile,
            '-U', 'ALLOW_UNINDEXED_BAM']
    
    return args


def table_recalibration(reference_genome,bam_file_name,recal_file_name, outfile, use_mem, path_to_gatk):
    '''
    Generates a new bam file with recalibrated quality scores
    
    java -Xmx16g -jar /exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/GenomeAnalysisTK.jar 
    -l INFO -R /exds/projects/alignment_indexes/gatk/hg19/hg19.fa -I s_3_12_sequence.hg19.aln2.rmdup.sorted.bam 
    -T TableRecalibration -o s_3_12_sequence.aln2.rmdup.sorted.recal.bam 
    -recalFile s_3_12_sequence.hg19.aln2.rmdup.sorted.before.recal_data.csv
    The original function in gatk does not support multiple threads yet
    '''
    gatk_command=path_to_gatk+'GenomeAnalysisTK.jar'
    #outfile=bam_file_name.replace('.bam','.recal.bam')
    
    args=['java','-Xmx'+str(use_mem)+'m', '-jar',gatk_command,
            '-l', 'INFO','-R',reference_genome, '-I', bam_file_name, '-T','TableRecalibration',
            '-o',outfile, '-recalFile',recal_file_name,
            '-U', 'ALLOW_UNINDEXED_BAM']

    return args


def analyze_covariates(recal_file_name,outputdir, resources_folder, rscipt_path, use_mem, path_to_gatk):
    '''
    In this script is important to specify the path to resources in which you have the Rscripts for making the plots and the -R scripts
    -resources /exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/resources (resources_folder)
    -Rscript /exds/sw/local/R-project/bin/Rscript (rscipt_path)
    -ignoreQ 5
    
    java -Xmx16g -jar /exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/AnalyzeCovariates.jar 
    -recalFile s_3_12_sequence.hg19.aln2.rmdup.sorted.recal_data.csv 
    -outputDir /exds/users/oabalbin/projects/snps/exomes/aM18/recal_analysis/before/ 
    -resources /exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/resources/ 
    -Rscript /exds/sw/local/R-project/bin/Rscript
    '''
    
    gatk_command=path_to_gatk+'AnalyzeCovariates.jar'
    args=['java','-Xmx'+str(use_mem)+'m', '-jar',gatk_command,
            '-recalFile',recal_file_name,'-outputDir',outputdir,'-resources',resources_folder,'-Rscript',rscipt_path]
    return args


def check_create_dir(root_path, dir_name):
    '''
        if not os.path.isdir(full_path_name):
            os.mkdir( full_path_name )
    '''
    if not os.path.isdir(root_path):
            os.mkdir( root_path )
            
    subfolder=os.path.join(root_path,dir_name)
    
    if not os.path.isdir(subfolder):
            os.mkdir( subfolder )
 
    return subfolder



def main_baseQ_recalibration(jname, configrun, thislane, bam_file_name, core_mem, num_cores):
    '''
    This is a wrap for the base quality recalibration using  gatk (Broad).
    Input: reference genome, vcf dbsnp database, indexed original bam file, 
        memory available, num of cores to use, analysis output directory
        path to gatk, resources folder, and RScript
    
    Returns: recalibrate bam file path, Writes recalibrated bam file and 
        Count covariate tables
        
    gatk quality recalibration steps
    CountCovariates, TableRecalibrate, 
    samtools index on the recalibrated bam file, CountCovariates
    analyze covariates before and after recalibration
    Most of the parameters in this function should go into a dictionary structure.
    '''
    # Parameters:
    from exome.jobs.job_runner import run_local
    jobrunfunc=run_local

    extra_mem, num_cores = configrun.gatk_use_mem, configrun.gatk_num_cores
    path_to_gatk, path_to_picard, path_to_sam = configrun.gatk_path, configrun.picard_path, \
                                                configrun.samtools_path
    genomes = configrun.genomes['human']
    ref_genome, ref_dbsnp, indeldb_file = genomes.gatk_ref_genome, genomes.snpdb, \
                                         genomes.indeldb        
    recal_analysis_outputdir=configrun.gatk_recal_anal_dir
    temp_dir =  configrun.gatk_temp_dir
    resources_folder,rscipt_path = configrun.gatk_resource_path, configrun.rscript_bin
    log_dir=thislane.gatk_dir
    gatk_covariates_analysis = configrun.gatk_covariates_analysis
    #bam_file_name = thislane.markdup_bam_file
    
    # Package 1
    # Generate count covariates table for recalibration   
    pre_recal_table_file = thislane.pre_recal_table_file
    jobidcv=-1
    if up_to_date(pre_recal_table_file,pre_recal_table_file) and up_to_date_job(pre_recal_table_file):
        logging.info("[SKIPPED] Count covariates step. File %s file is up to date" % (pre_recal_table_file))
        jobidcv=0
    else:
        logging.info("Starting Count covariates step. File %s" % (pre_recal_table_file))

        command = count_covariates(ref_genome, ref_dbsnp, bam_file_name,
                                   pre_recal_table_file, core_mem, num_cores, path_to_gatk)
        
        jobidcv = jobrunfunc(jname, command, num_cores, cwd=log_dir, stdout='count_covariates.log')
        if jobidcv == 0:
                job_done(pre_recal_table_file)

    # Recalibrate original bam file
    # Package 1 starts: into one script
    
    if jobidcv != 0:
        raise OSError("command: '%s' returned error code %s" % (' '.join(command), jobidcv))
    else:
        recal_bam_file_name=thislane.recal_bam_file_name
        recal_bam_indexed_name=thislane.recal_bam_indexed_name
        if up_to_date(recal_bam_indexed_name,recal_bam_indexed_name) and up_to_date_job(recal_bam_indexed_name):
            logging.info("[SKIPPED] Create Table recalibration step. File %s file is up to date" % (recal_bam_indexed_name))
            jobidtr=0
        else:
            logging.info("Starting Create Table recalibration step %s" % (recal_bam_file_name))
    
            command = table_recalibration(ref_genome, bam_file_name,
                                                      pre_recal_table_file, recal_bam_file_name, 
                                                      extra_mem, path_to_gatk)
            
            jobidtr = jobrunfunc(jname, command, num_cores, cwd=log_dir, stdout='Tale_recalibration.log')
            
            if jobidtr==0:
                command  = mypicard.sortIndexSam(recal_bam_file_name, recal_bam_indexed_name, 
                                                 extra_mem, path_to_picard)
                
                jobidib =jobrunfunc(jname, command, num_cores, cwd=log_dir, stdout='SortRecalibrateFile.log')
                if jobidib != 0:
                    job_done(recal_bam_indexed_name)

            else:
                jobidib=-1

                
    # Package 1 ends:
    
    # Generate count covariates table for recalibrated bam file
    # Here maybe using more cores
    if jobidib != 0:
        raise OSError("command: '%s' returned error code %s" % (' '.join(command), jobidib))
    elif gatk_covariates_analysis:
        post_recal_table_file =thislane.post_recal_table_file
        if up_to_date(post_recal_table_file,post_recal_table_file) and up_to_date_job(post_recal_table_file):
            logging.info("[SKIPPED] Covariates Analysis step. File %s file is up to date" % (post_recal_table_file))
            jobidcv2=0
        else:
            logging.info("Starting Covariate Analysis step %s" % (post_recal_table_file))

            command = count_covariates(ref_genome, ref_dbsnp, recal_bam_indexed_name, 
                                       post_recal_table_file, core_mem, num_cores, path_to_gatk)

            jobidcv2=jobrunfunc(jname, command, num_cores, cwd=log_dir, stdout='CovariateAnalysis.log')
            
            if jobidcv2!=0:
                raise OSError("command: '%s' returned error code %s" % (' '.join(command), jobidib))
            else:
        
                # Analyze covariates before and after recalibration
                #Creating folder should be done for another job during the setup
                    
                before_recal_results_path = check_create_dir(recal_analysis_outputdir,'before')
                after_recal_results_path = check_create_dir(recal_analysis_outputdir,'after')    
                # use the retcode to print to a logfile
                commandB = analyze_covariates(pre_recal_table_file, before_recal_results_path, 
                                             resources_folder, rscipt_path, extra_mem, path_to_gatk)
                #
                commandA = analyze_covariates(post_recal_table_file, after_recal_results_path, 
                                              resources_folder, rscipt_path, extra_mem, path_to_gatk)
                
                # Change this line
                command = commandA+'\n'+commandB
                jobidab=jobrunfunc(jname, command, num_cores, cwd=log_dir, stdout='CovariateAnalysis2.log')
                if jobidab==0 or jobidib==0:
                    return JOB_SUCCESS
                else:
                    JOB_ERROR
        
    elif jobidib == 0:
        return JOB_SUCCESS
    else:
        JOB_ERROR


def main():
    logging.basicConfig(level=logging.DEBUG,
                        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
    
    
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-r", "--config_file", dest="config_file",
                            help="configuration run object")
    '''
    optionparser.add_option("-a", "--analysis_file", dest="analysis_file",
                            help="file with experiment configuration") 
    '''
    optionparser.add_option("-j", "--jname", dest="jname",
                            help="jobname") 
    optionparser.add_option("-l", "--lane_file", dest="lane_file",
                            help="lane object with information about the lane to be processed") 
    optionparser.add_option("-d", "--lane_root_dir", dest="lane_root_dir",
                            help="lane object with information about the lane to be processed") 
    
    optionparser.add_option("-i", "--bamfile", dest="bamfile",
                            help="bamfile for base quality recalibration, it needs to be sorted and indexed.") 

    optionparser.add_option("-c", "--ncores", dest="ncores",
                            help="ncores to use in the analysis") 
    optionparser.add_option("-m", "--core_mem", dest="core_mem",
                            help="Amount of memory allocated for this process.") 
    

    (options, args) = optionparser.parse_args()
    
    configrun = ExomePipelineConfig()
    configrun.from_xml(options.config_file)
    thlane  = ExomeLaneConfig()
    #job.from_xml(lane_file, config.output_dir)
    thlane.from_xml(options.lane_file, options.lane_root_dir)

    return main_baseQ_recalibration(options.jname, configrun, thlane, 
                                    options.bamfile, options.core_mem, options.ncores)

if __name__ == '__main__': 
    import sys
    sys.exit(main())
    
    
    




    
    
    
    
